import numpy as np
import matplotlib.pyplot as mp
import partie1_methodes as p1

### Test global ###

def test_meth():
	"verifie l'exactitude de meth_epsilon (et donc de meth_n_step) dans plusieurs cas"
	print "Tests de correction des methodes de resolution...\n"
	
	# y = e^x
	print "Equation : (y\'=y, y(0)=1) -> Solution : y = e^x"
	eq = [[0.,1.], lambda y,t: y]
	tf = 3.
	epsilon = 0.001
	erreur = 0.1
	
	Y = p1.meth_epsilon(eq,tf,epsilon,p1.step_euler)
	X = np.arange(0, tf, tf/len(Y))
	Yref = np.exp(X)
	assert((abs(Y - Yref) < erreur).all)
	print "	Euler : ok"
	
	Y = p1.meth_epsilon(eq,tf,epsilon,p1.step_ptmilieu)
	X = np.arange(0, tf, tf/len(Y))
	Yref = np.exp(X)
	assert((abs(Y - Yref) < erreur).all)
	print "	Point milieu : ok"
	
	Y = p1.meth_epsilon(eq,tf,epsilon,p1.step_heun)
	X = np.arange(0, tf, tf/len(Y))
	Yref = np.exp(X)
	assert((abs(Y - Yref) < erreur).all)
	print "	Heun : ok"
	
	Y = p1.meth_epsilon(eq,tf,epsilon,p1.step_RK4)
	X = np.arange(0, tf, tf/len(Y))
	Yref = np.exp(X)
	assert((abs(Y - Yref) < erreur).all)
	print "	RK4 : ok \n"
	
	# y = e^arctan(x)
	print "Equation : (y\'(t)=1/(1+t^2), y(0)=1) -> Solution : y = e^arctan(x)"
	eq = [[0.,1.], lambda y,t: y/(1+t**2)]
	tf = 3.
	epsilon = 0.001
	erreur = 0.1
	
	Y = p1.meth_epsilon(eq,tf,epsilon,p1.step_euler)
	X = np.arange(0, tf, tf/len(Y))
	Yref = np.exp(X)
	assert((abs(Y - Yref) < erreur).all)
	print "	Euler : ok"
	
	Y = p1.meth_epsilon(eq,tf,epsilon,p1.step_ptmilieu)
	X = np.arange(0, tf, tf/len(Y))
	Yref = np.exp(X)
	assert((abs(Y - Yref) < erreur).all)
	print "	Point milieu : ok"
	
	Y = p1.meth_epsilon(eq,tf,epsilon,p1.step_heun)
	X = np.arange(0, tf, tf/len(Y))
	Yref = np.exp(X)
	assert((abs(Y - Yref) < erreur).all)
	print "	Heun : ok"
	
	Y = p1.meth_epsilon(eq,tf,epsilon,p1.step_RK4)
	X = np.arange(0, tf, tf/len(Y))
	Yref = np.exp(X)
	assert((abs(Y - Yref) < erreur).all)
	print "	RK4 : ok\n"

### Applications graphiques des diverses methodes ###

def plot_expo1():
	"graphique des approximations de e^x obtenues avec un pas de 1"
	Yg = [[0.,1.], lambda y,t: y]
	mp.xlim(0,3)
	mp.ylim(1,5)
	x = np.arange(0, 5, 1)
	mp.plot(x,np.exp(x),label='y=e^x')
	mp.plot(x,p1.meth_n_step(Yg, 4, 1, p1.step_euler), label='Euler')
	mp.plot(x,p1.meth_n_step(Yg, 4, 1, p1.step_ptmilieu), label='Pt milieu')
	mp.plot(x,p1.meth_n_step(Yg, 4, 1, p1.step_heun), label='Heun')
	mp.plot(x,p1.meth_n_step(Yg, 4, 1, p1.step_RK4), label='RK4')
	mp.title('Resolution de (y\'=y, y(0)=1) avec h=1')
	mp.legend()
	mp.show()

def plot_expo2():
	"graphique des approximations de e^x obtenues avec un pas de 0.1"
	Yg = [[0.,1.], lambda y,t: y]
	mp.xlim(0,5)
	mp.ylim(1,30)
	x = np.arange(0, 10, 0.1)
	mp.plot(x,np.exp(x),label='y=e^x')
	mp.plot(x,p1.meth_n_step(Yg, 99, 0.1, p1.step_euler), label='Euler')
	mp.plot(x,p1.meth_n_step(Yg, 99, 0.1, p1.step_ptmilieu), label='Pt milieu')
	mp.plot(x,p1.meth_n_step(Yg, 99, 0.1, p1.step_heun), label='Heun')
	mp.plot(x,p1.meth_n_step(Yg, 99, 0.1, p1.step_RK4), label='RK4')
	mp.title('Resolution de (y\'=y, y(0)=1) avec h=0.1')
	mp.legend()
	mp.show()
	
def plot_arctan_expo():
	"graphique des approximations de e^arctan(x) obtenues avec un pas de 1"
	xmax = 10
	h=1
	N=xmax - 1
	eq = [[0.,1.], lambda y,t: y/(1+t**2)]
	mp.xlim(0,xmax-1)
	mp.ylim(1,6.2)
	x = np.arange(0, xmax, h)
	mp.plot(x,np.exp(np.arctan(x)),label='y=e^arctan(x)')
	mp.plot(x,p1.meth_n_step(eq, N, h, p1.step_euler), label='Euler')
	mp.plot(x,p1.meth_n_step(eq, N, h, p1.step_ptmilieu), label='Pt milieu')
	mp.plot(x,p1.meth_n_step(eq, N, h, p1.step_heun), label='Heun')
	mp.plot(x,p1.meth_n_step(eq, N, h, p1.step_RK4), label='RK4')
	mp.title('Resolution de (y\'(t)=1/(1+t^2), y(0)=1) avec h=1')
	mp.legend()
	mp.show()

def plot_equation_cercle1():
	"Graphique approximations de cercle obtenues pour Euler avec differents pas"
	
	t0 = 0.
	y0 = [1., 0.]
	eq = [(t0,y0),lambda X,t:np.array([-X[1],X[0]])]
	times = 2*np.pi
	
	N = 5
	h = times/N
	res = p1.meth_n_step(eq,N, h,p1.step_euler)
	tabres0 = map(lambda t : t[0], res)
	tabres1 = map(lambda t : t[1], res)
	mp.plot(tabres1, tabres0, linewidth=1.0, label='N=5')
	
	N = 10
	h = times/N
	res = p1.meth_n_step(eq,N, h,p1.step_euler)
	tabres0 = map(lambda t : t[0], res)
	tabres1 = map(lambda t : t[1], res)
	mp.plot(tabres1, tabres0, linewidth=1.0, label='N=10')
	
	N = 20
	h = times/N
	res = p1.meth_n_step(eq,N, h,p1.step_euler)
	tabres0 = map(lambda t : t[0], res)
	tabres1 = map(lambda t : t[1], res)
	mp.plot(tabres1, tabres0, linewidth=1.0, label='N=20')
	
	N = 100
	h = times/N
	res = p1.meth_n_step(eq,N, h,p1.step_RK4)
	tabres0 = map(lambda t : t[0], res)
	tabres1 = map(lambda t : t[1], res)
	mp.plot(tabres1, tabres0, linewidth=1.0, label='reference')
	
	mp.title('Approximation du cercle avec Euler (h=2pi/N)')
	mp.legend()
	mp.show()

def plot_equation_cercle2():
	"Graphique des approximations de cercle obtenues pour le Pt milieu, Heun et RK4"
	N = 10
	times = 2*np.pi
	h = times/N

	t0 = 0.
	y0 = [1., 0.]
	eq = [(t0,y0),lambda X,t:np.array([-X[1],X[0]])]
	
	#Euler
	res = p1.meth_n_step(eq,N, h,p1.step_euler)
	tabres0 = map(lambda t : t[0], res)
	tabres1 = map(lambda t : t[1], res)
	mp.plot(tabres1, tabres0, linewidth=1.0, label='Euler')

	#Point milieu
	res = p1.meth_n_step(eq,N, h,p1.step_ptmilieu)
	tabres0 = map(lambda t : t[0], res)
	tabres1 = map(lambda t : t[1], res)
	mp.plot(tabres1, tabres0, linewidth=1.0, label='Pt milieu')
	
	#Heun
	res = p1.meth_n_step(eq,N, h,p1.step_heun)
	tabres0 = map(lambda t : t[0], res)
	tabres1 = map(lambda t : t[1], res)
	mp.plot(tabres1, tabres0, linewidth=1.0, label='Heun')

	#RK4
	res = p1.meth_n_step(eq,N, h,p1.step_RK4)
	tabres0 = map(lambda t : t[0], res)
	tabres1 = map(lambda t : t[1], res)
	mp.plot(tabres1, tabres0, linewidth=1.0, label='RK4')

	mp.title('Equation de cercle (N=10,h=2*pi/N)')
	mp.legend()
	mp.show()

def plot_expo_meth_epsilon():
	Yg = [[0.,1.], lambda y,t: y]
	tf = 3.
	mp.xlim(0,2)
	mp.ylim(1,4)
	x = np.arange(0, tf, 0.01)
	mp.plot(x,np.exp(x),label='y=e^x')
	
	epsilon = 2
	y = p1.meth_epsilon(Yg,tf,epsilon,p1.step_euler)
	x = np.arange(0, tf, tf/len(y))
	mp.plot(x,y, label='eps=2')
	
	epsilon = 1
	y = p1.meth_epsilon(Yg,tf,epsilon,p1.step_euler)
	x = np.arange(0, tf, tf/len(y))
	mp.plot(x,y, label='eps=1')
	
	epsilon = 0.1
	y = p1.meth_epsilon(Yg,tf,epsilon,p1.step_euler)
	x = np.arange(0, tf, tf/len(y))
	mp.plot(x,y, label='eps=0.1')
	
	
	mp.title('Resolution de (y\'=y, y(0)=1) avec epsilon variable')
	mp.legend()
	mp.show()
	
def plot_champs_vecteurs():
	eq = [[-2,np.exp(-4)],lambda t,y: -2*t*y]
	p1.champs_vecteurs(eq)
	
	eq = [[0,1],lambda t,y: y]
	p1.champs_vecteurs(eq,0,4)


if __name__ == "__main__":
	test_meth()
	plot_expo1()
	plot_expo2() 
	plot_arctan_expo()
	plot_equation_cercle1()
	plot_equation_cercle2()
	plot_champs_vecteurs()
